#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2017-09-26 16:40:15
# @Author  : WUQINGLONG
# @Contact : 1603763721@qq.com

# =======================================
# 优美胜于丑陋（Python 以编写优美的代码为目标）
# 明了胜于晦涩（优美的代码应当是明了的，命名规范，风格相似）
# 简洁胜于复杂（优美的代码应当是简洁的，不要有复杂的内部实现）
# 复杂胜于凌乱（如果复杂不可避免，那代码间也不能有难懂的关系，要保持接口简洁）
# 扁平胜于嵌套（优美的代码应当是扁平的，不能有太多的嵌套）
# 间隔胜于紧凑（优美的代码有适当的间隔，不要奢望一行代码解决问题）
# 可读性很重要（优美的代码是可读的）
# =======================================

"""
功能描述：
    1.aes加密解密工具类

"""

import base64
import hashlib
import random
import sys
from math import ceil

from Crypto.Cipher import AES

# 默认秘钥
KEY = 'gMIxtfBaCxmcoG5E'
# 填充值
PADDING = b'\0'
# 填充值
IV = PADDING * 16
# Python版本判断
VERSION2 = (sys.version_info[0] == 2)


def _padding(data='', size=16):
    """字符串扩容，使长度为16的倍数"""
    if not VERSION2:
        data = data.encode('utf-8')
    return data + PADDING * (size - len(data) % size)


def _generator(key=KEY):
    """根据指定的Key生成密码器"""
    return AES.new(key, AES.MODE_CBC, IV)


def encrypt(data=None, key=KEY, key_size=0):
    """
    加密数据
        1.key为秘钥
        2.key_size为自省秘钥长度，设置后，随机生成指定长度的秘钥，并追加在返回的加密串之后
    """
    if data is None:
        return
    # 设置自省秘钥时，需要自动生成秘钥
    if key_size:
        key = random_key(key_size)
    crypt = _generator(key).encrypt(_padding(data))
    if key_size:
        return (base64.b64encode(crypt) + key).decode('ascii')
    else:
        return base64.b64encode(crypt).decode('ascii')


def decrypt(data=None, key=KEY, key_size=0):
    """
    解密数据
        1.key为秘钥
        2.key_size为自省秘钥长度，设置后，从加密的字符串最后截取指定长度作为秘钥
    """
    if data is None:
        return
    # 设置自省秘钥长度时，需要分割解密字符串和秘钥
    if key_size:
        key_size = int(key_size)
        key = data[-key_size:]
        data = data[0:-key_size]
    crypt = base64.b64decode(data)
    recovery = _generator(key).decrypt(crypt)
    return recovery.rstrip(PADDING).decode('utf-8')


def random_key(size=16):
    """随机生成一个size长度的加密key"""
    m5 = hashlib.md5()
    m5.update(str(random.random()))
    key = m5.hexdigest()
    length_key = len(key)
    # md5只有32位，如果size超过该长度需要扩容
    if length_key < size:
        keys = list(key * int(ceil(size / length_key)))
        # 扩容后打乱顺序
        random.shuffle(keys)
        key = ''.join(keys)
    return key[0:size]


if __name__ == '__main__':
    txt = 'hello world!你好，中国'
    encrypt_txt = encrypt(txt)
    print('原始的字符串：' + '%s' % txt)
    print(u'加密的字符串：' + '%s' % encrypt_txt)
    print(u'解密的字符串：' + '%s' % decrypt(encrypt_txt))
    txt = '{"ok":404}'
    encrypt_txt = encrypt(txt)
    print('原始的字符串：%s' % txt)
    print(u'加密的字符串：' + '%s' % encrypt_txt)
    print(u'解密的字符串：' + '%s' % decrypt(encrypt_txt))
